"""Main controller for GHOSTCREW application."""

import asyncio
import traceback
from datetime import datetime
from typing import Optional, List, Dict, Any
from colorama import Fore, Style

from config.constants import (
    ASCII_TITLE, VERSION, WELCOME_MESSAGE, EXIT_MESSAGE, SEPARATOR,
    KB_PROMPT, MCP_PROMPT, ERROR_NO_WORKFLOWS, ERROR_NO_REPORTING,
    DEFAULT_KNOWLEDGE_BASE_PATH
)
from config.app_config import app_config
from core.agent_runner import agent_runner
from core.agent_mode_controller import AgentModeController
from tools.mcp_manager import MCPManager
from ui.menu_system import MenuSystem
from ui.conversation_manager import ConversationManager
from workflows.workflow_engine import WorkflowEngine
from rag.knowledge_base import Kb


class PentestAgent:
    """Main application controller for GHOSTCREW."""
    
    def __init__(self, MCPServerStdio=None, MCPServerSse=None):
        """
        Initialize the pentest agent controller.
        
        Args:
            MCPServerStdio: MCP server stdio class
            MCPServerSse: MCP server SSE class
        """
        self.app_config = app_config
        self.agent_runner = agent_runner
        self.mcp_manager = MCPManager(MCPServerStdio, MCPServerSse)
        self.menu_system = MenuSystem()
        self.conversation_manager = ConversationManager()
        self.workflow_engine = WorkflowEngine()
        self.kb_instance = None
        self.reporting_available = self._check_reporting_available()
        self.agent_mode_controller = None  # Will be initialized when needed
    
    @staticmethod
    def _check_reporting_available() -> bool:
        """Check if reporting module is available."""
        try:
            from reporting.generators import generate_report_from_workflow
            return True
        except ImportError:
            print(ERROR_NO_REPORTING)
            return False
    
    def display_welcome(self) -> None:
        """Display welcome message and ASCII art."""
        print(ASCII_TITLE)
        print(f"{Fore.WHITE}GHOSTCREW v{VERSION}{Style.RESET_ALL}")
        print(WELCOME_MESSAGE)
        print(EXIT_MESSAGE)
        print(f"{SEPARATOR}\n")
    
    def setup_knowledge_base(self) -> None:
        """Setup knowledge base if requested by user."""
        use_kb_input = input(KB_PROMPT).strip().lower()
        if use_kb_input == 'yes':
            try:
                self.kb_instance = Kb(DEFAULT_KNOWLEDGE_BASE_PATH)
                print(f"{Fore.GREEN}Knowledge base loaded successfully!{Style.RESET_ALL}")
            except Exception as e:
                print(f"{Fore.RED}Failed to load knowledge base: {e}{Style.RESET_ALL}")
                self.kb_instance = None
    
    async def setup_mcp_tools(self) -> tuple:
        """Setup MCP tools and return server instances."""
        use_mcp_input = input(MCP_PROMPT).strip().lower()
        return await self.mcp_manager.setup_mcp_tools(use_mcp_input == 'yes')
    
    async def run_interactive_mode(self, connected_servers: List) -> None:
        """Run interactive chat mode."""
        self.menu_system.display_interactive_mode_intro()
        
        while True:
            user_query = self.menu_system.get_user_input()
            
            # Handle special commands
            if user_query.lower() in ["quit", "exit"]:
                self.menu_system.display_exit_message()
                return True  # Signal to exit the entire application
            
            if user_query.lower() == "menu":
                break  # Return to main menu
            
            # Handle empty input
            if not user_query:
                self.menu_system.display_no_query_message()
                continue
            
            # Handle multi-line mode request
            if user_query.lower() == "multi":
                user_query = self.menu_system.get_multi_line_input()
                if not user_query:
                    continue
            
            # Process the query
            await self._process_user_query(user_query, connected_servers)
            
            self.menu_system.display_ready_message()
        
        return False  # Don't exit application
    
    async def _process_user_query(self, query: str, connected_servers: List) -> None:
        """Process a user query through the agent."""
        # Add dialogue to history
        self.conversation_manager.add_dialogue(query)
        
        # Run the agent
        result = await agent_runner.run_agent(
            query,
            connected_servers,
            history=self.conversation_manager.get_history(),
            streaming=True,
            kb_instance=self.kb_instance
        )
        
        # Update the response in history
        if result and hasattr(result, "final_output"):
            self.conversation_manager.update_last_response(result.final_output)
    
    async def run_automated_mode(self, connected_servers: List) -> None:
        """Run workflows mode."""
        if not self.workflow_engine.is_available():
            print(ERROR_NO_WORKFLOWS)
            self.menu_system.press_enter_to_continue()
            return
        
        if not connected_servers:
            self.menu_system.display_workflow_requirements_message()
            return
        
        while True:
            workflow_list = self.workflow_engine.show_automated_menu()
            if not workflow_list:
                break
            
            try:
                choice = input(f"\n{Fore.GREEN}Select workflow (1-{len(workflow_list)+1}): {Style.RESET_ALL}").strip()
                
                if not choice.isdigit():
                    self.menu_system.display_invalid_input()
                    continue
                
                choice = int(choice)
                
                if choice == len(workflow_list) + 1:
                    # Back to main menu
                    break
                
                if 1 <= choice <= len(workflow_list):
                    await self._execute_workflow(workflow_list[choice-1], connected_servers)
                else:
                    self.menu_system.display_invalid_choice()
            
            except ValueError:
                self.menu_system.display_invalid_input()
            except KeyboardInterrupt:
                self.menu_system.display_operation_cancelled()
                break
    
    async def run_agent_mode(self, connected_servers: List) -> None:
        """Run autonomous agent mode with PTT."""
        if not connected_servers:
            self.menu_system.display_agent_mode_requirements_message()
            return
        
        # Display introduction
        self.menu_system.display_agent_mode_intro()
        
        # Get agent mode parameters
        params = self.menu_system.get_agent_mode_params()
        if not params:
            return
        
        # Initialize agent mode controller
        self.agent_mode_controller = AgentModeController(
            self.mcp_manager,
            self.conversation_manager,
            self.kb_instance
        )
        
        try:
            # Initialize agent mode
            init_success = await self.agent_mode_controller.initialize_agent_mode(
                goal=params['goal'],
                target=params['target'],
                constraints=params['constraints'],
                connected_servers=connected_servers,
                run_agent_func=agent_runner.run_agent
            )
            
            if init_success:
                # Run the autonomous loop
                await self.agent_mode_controller.run_autonomous_loop()
                
                # Handle post-execution options
                await self._handle_agent_mode_completion()
            else:
                print(f"{Fore.RED}Failed to initialize agent mode.{Style.RESET_ALL}")
                self.menu_system.press_enter_to_continue()
        
        except KeyboardInterrupt:
            print(f"\n{Fore.YELLOW}Agent mode interrupted by user.{Style.RESET_ALL}")
        except Exception as e:
            print(f"{Fore.RED}Error in agent mode: {e}{Style.RESET_ALL}")
            traceback.print_exc()
        finally:
            self.menu_system.press_enter_to_continue()
    
    async def _handle_agent_mode_completion(self) -> None:
        """Handle post-execution options for agent mode."""
        # Ask if user wants to generate a report
        if self.reporting_available and self.menu_system.ask_generate_report():
            try:
                # Generate report from PTT
                from reporting.generators import generate_report_from_ptt
                
                ptt = self.agent_mode_controller.get_ptt_for_reporting()
                report_path = await generate_report_from_ptt(
                    ptt,
                    self.conversation_manager.get_history(),
                    run_agent_func=agent_runner.run_agent,
                    connected_servers=self.mcp_manager.connected_servers if hasattr(self.mcp_manager, 'connected_servers') else [],
                    kb_instance=self.kb_instance
                )
                
                if report_path:
                    self.menu_system.display_report_generated(report_path)
                else:
                    print(f"{Fore.YELLOW}Report generation returned no path.{Style.RESET_ALL}")
            
            except ImportError:
                # Fallback if PTT report generation not available
                print(f"{Fore.YELLOW}PTT report generation not available. Saving raw data...{Style.RESET_ALL}")
                self._save_agent_mode_data()
            except Exception as e:
                self.menu_system.display_report_error(e)
                self._save_agent_mode_data()
        
        # Ask about saving raw data
        elif self.menu_system.ask_save_raw_history():
            self._save_agent_mode_data()
    
    def _save_agent_mode_data(self) -> None:
        """Save agent mode execution data."""
        try:
            import os
            import json
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            
            # Create reports directory if it doesn't exist
            os.makedirs("reports", exist_ok=True)
            
            # Save conversation history
            history_file = f"reports/agent_mode_history_{timestamp}.json"
            with open(history_file, 'w', encoding='utf-8') as f:
                json.dump(self.conversation_manager.get_history(), f, indent=2)
            
            # Save PTT state
            if self.agent_mode_controller:
                ptt_file = f"reports/agent_mode_ptt_{timestamp}.json"
                with open(ptt_file, 'w', encoding='utf-8') as f:
                    f.write(self.agent_mode_controller.tree_manager.to_json())
            
            print(f"{Fore.GREEN}Agent mode data saved to reports/ directory{Style.RESET_ALL}")
        
        except Exception as e:
            print(f"{Fore.RED}Failed to save agent mode data: {e}{Style.RESET_ALL}")
    
    async def _execute_workflow(self, workflow_info: tuple, connected_servers: List) -> None:
        """Execute a selected workflow."""
        workflow_key, workflow_name = workflow_info
        workflow = self.workflow_engine.get_workflow(workflow_key)
        
        if not workflow:
            print(f"{Fore.RED}Error loading workflow.{Style.RESET_ALL}")
            return
        
        target = self.menu_system.get_workflow_target()
        if not target:
            return
        
        if not self.menu_system.confirm_workflow_execution(workflow['name'], target):
            self.menu_system.display_workflow_cancelled()
            return
        
        # Store initial workflow data
        workflow_start_time = datetime.now()
        initial_history_length = self.conversation_manager.get_dialogue_count()
        
        # Execute the workflow
        await self.workflow_engine.run_automated_workflow(
            workflow,
            target,
            connected_servers,
            self.conversation_manager.get_history(),
            self.kb_instance,
            agent_runner.run_agent
        )
        
        self.menu_system.display_workflow_completed()
        
        # Handle report generation
        if self.reporting_available:
            await self._handle_report_generation(
                workflow,
                workflow_key,
                target,
                workflow_start_time,
                initial_history_length,
                connected_servers
            )
        else:
            print(f"\n{Fore.YELLOW}Reporting not available.{Style.RESET_ALL}")
        
        self.menu_system.press_enter_to_continue()
    
    async def _handle_report_generation(
        self,
        workflow: Dict,
        workflow_key: str,
        target: str,
        workflow_start_time: datetime,
        initial_history_length: int,
        connected_servers: List
    ) -> None:
        """Handle report generation after workflow completion."""
        if not self.menu_system.ask_generate_report():
            return
        
        save_raw_history = self.menu_system.ask_save_raw_history()
        
        try:
            from reporting.generators import generate_report_from_workflow
            
            # Prepare report data
            workflow_conversation = self.conversation_manager.get_workflow_conversation(initial_history_length)
            
            report_data = {
                'workflow_name': workflow['name'],
                'workflow_key': workflow_key,
                'target': target,
                'timestamp': workflow_start_time,
                'conversation_history': workflow_conversation,
                'tools_used': MCPManager.get_available_tools(connected_servers)
            }
            
            # Generate professional report
            print(f"\n{Fore.CYAN}Generating report...{Style.RESET_ALL}")
            report_path = await generate_report_from_workflow(
                report_data,
                agent_runner.run_agent,
                connected_servers,
                self.kb_instance,
                save_raw_history
            )
            
            self.menu_system.display_report_generated(report_path)
            
        except Exception as e:
            self.menu_system.display_report_error(e)
    
    async def run(self) -> None:
        """Main application run method."""
        self.display_welcome()
        self.setup_knowledge_base()
        
        try:
            # Setup MCP tools
            mcp_server_instances, connected_servers = await self.setup_mcp_tools()
            
            # Check if we need to restart (e.g., after configuring new tools)
            if mcp_server_instances and not connected_servers:
                return
            
            # Main application loop
            while True:
                self.menu_system.display_main_menu(
                    self.workflow_engine.is_available(),
                    bool(connected_servers)
                )
                
                menu_choice = self.menu_system.get_menu_choice()
                
                if menu_choice == "1":
                    # Interactive mode
                    should_exit = await self.run_interactive_mode(connected_servers)
                    if should_exit:
                        break
                
                elif menu_choice == "2":
                    # Automated mode
                    await self.run_automated_mode(connected_servers)
                
                elif menu_choice == "3":
                    # Agent mode
                    await self.run_agent_mode(connected_servers)
                
                elif menu_choice == "4":
                    # Exit
                    self.menu_system.display_exit_message()
                    break
                
                else:
                    self.menu_system.display_invalid_choice()
        
        except KeyboardInterrupt:
            print(f"\n{Fore.YELLOW}Program interrupted by user, exiting...{Style.RESET_ALL}")
        except Exception as e:
            print(f"{Fore.RED}Error during program execution: {e}{Style.RESET_ALL}")
            traceback.print_exc()
        finally:
            # Cleanup MCP servers
            await self.mcp_manager.cleanup_servers()
            
            # Close any remaining asyncio transports
            await self._cleanup_asyncio_resources()
            
            print(f"{Fore.GREEN}Program ended.{Style.RESET_ALL}")
    
    async def _cleanup_asyncio_resources(self) -> None:
        """Clean up asyncio resources."""
        try:
            # Get the event loop
            loop = asyncio.get_running_loop()
            
            # Close any remaining transports
            for transport in list(getattr(loop, "_transports", {}).values()):
                if hasattr(transport, "close"):
                    try:
                        transport.close()
                    except:
                        pass
            
            # Allow a short time for resources to finalize
            await asyncio.sleep(0.1)
        except:
            pass  # Ignore any errors in the final cleanup 